'''
Created on 29.01.2012

@author: myaut
'''

import subprocess
import threading
import os
import signal
import time

from Queue import Queue

import salsa
from salsa import trace_exec_flow, do_trace, SALSAException

from salsa.util.io import IOFeederReply, IODeferredReply
from salsa.util import LoggerSeverity
from salsa.util.rpc import RPCProtocol, RPCError

TASK_INTR_WAIT = 3
TASK_RESTART_ATTEMPTS = 3

task_start_lock = threading.Lock()

class TaskMgrException(SALSAException):
    pass

class TaskWaiterEvent:
    '''Adapter for threading.Event that counts count times before wakes up threads
    Required for waiting on multiple tasks, used in IODefferedReply'''
    def __init__(self, count):
        self.count = count
        self.lock = threading.Lock() 
        self.event = threading.Event()
    
    def wait(self, timeout = None):
        self.event.wait(timeout)
    
    def is_set(self):
        return self.count == 0
    
    def set(self):
        self.lock.acquire()
        
        try:
            if self.count > 0:
                self.count -= 1
            
            if self.count == 0:
                self.event.set()
        finally:
            self.lock.release()
            
        
class Application:
    py_debug = ['/usr/debug/bin/python']
    py_exec = ['/usr/bin/python']
    execs = {'python': py_exec, 'pydebug': py_debug}
    
    def __init__(self, name):
        self.name = name
        
    def set_params(self, params):
        self.params = params.copy()
        
    def set_binpath(self, binpath):
        self.pretty_binpath = binpath
        self.binpath = salsa.core.abs_bin_path(binpath)
        
        if not os.path.exists(self.binpath):
            raise TaskMgrException("Application '%s' doesn't exist" % binpath)
        
        if not os.access(self.binpath, os.R_OK | os.X_OK):
            raise TaskMgrException("Application '%s' is not executable/readable" % binpath)
    
    def set_exec(self, ename):
        if not ename:
            self.executor = []
            return
        
        if ename not in self.execs:
            raise TaskMgrException('Unknown application executor %s' % ename)
        
        self.executor = self.execs[ename]
         
    def __cmp__(self, other):
        return self.name.__cmp__(other.name)

class Task:
    _taskid = 0
    
    def __init__(self, app):
        self.app = app
        
        self.reset_taskid()
        
        self.proc = None
        self.stdout_file = None
        self.stderr_file = None
        
        self.finish_flag = False
        self.keep_alive = False
        self.restart_attempts = TASK_RESTART_ATTEMPTS
        
        self.env = {}
        self.args = {}
        
        self.waiter = None
        
        if trace_exec_flow: do_trace('Task %s is created' % self.taskid)
    
    def reset_taskid(self):                
        Task._taskid += 1
        
        self.tasknum = Task._taskid
        self.taskid = self.app.name + '-' + str(self.tasknum)
    
    def __str__(self):
        return 'taskid: %s pid: %d name: %s binpath: %s cwd: %s' % (self.taskid, 
                    self.proc.pid, self.app.name, self.app.pretty_binpath, self.cwd) 
    
    def get_pid(self):
        if self.proc:
            return self.proc.pid
        else:
            return 'X'
    
    def set_keep_alive(self):
        self.keep_alive = True
    
    def set_env(self, env):
        self.env = env
        self.env['TASK_ID'] = self.taskid
    
    def set_dirs(self, work_dir, output_dir):
        self.cwd = work_dir
        self.output_dir = output_dir
        
    def set_args(self, args):
        self.args = args
    
    def attach(self, waiter):
        if self.proc_finish_event.is_set():
            waiter.set()
        else:
            self.waiter = waiter
    
    def start(self, finish_handler):
        if not callable(finish_handler):
            raise TaskMgrException('finish_handler must be callable')
        
        self.proc_finish_event = threading.Event()
        self.finish_handler = finish_handler  
                
        self.thread = threading.Thread(target = self._run, name = 'TaskThread-%d' % self.tasknum)
        self.thread.start()
        
        if trace_exec_flow: do_trace('Task %s is started' % self.taskid)
    
    def _run(self):            
        while not self.finish_flag:
            if not self.keep_alive:
                # For not keep-alive tasks, stop them immediately after they finished
                self.finish_flag = True                
            
            self._run_task()
            
            if salsa.stopping:
                break 
            
            if self.keep_alive and not self.finish_flag:
                # Keep-alive tasks shouldn't die until finished explicitly by finish()
                # try to restart, if lack of restarts, stop SALSA
                self.restart_attempts -= 1
                
                salsa.core.logger.log(LoggerSeverity.LS_WARNING, 'Keep-alive task %s is dead, restart' % self.taskid)
                
                self.reset_taskid()
                
                if self.restart_attempts == 0:
                    salsa.core.fail('Task %s died %d times, but it is keep-alive task' % (self.taskid, TASK_RESTART_ATTEMPTS))
        
        if trace_exec_flow: do_trace('Task %s is finished' % self.taskid)
        
        self.finish_handler(self)
        self.proc_finish_event.set()
        
        # Notify waiter that we are finished
        if self.waiter:
            self.waiter.data[self.taskid] = self.proc_retcode
            self.waiter.wakeup()
        
    
    def _run_task(self):
        self.stdout_file = open(os.path.join(self.output_dir, self.taskid + '.out'), 'w')
        self.stderr_file = open(os.path.join(self.output_dir, self.taskid + '.err'), 'w')
        
        # Sometimes subprocess.Popen strangely waiting for GIL after forking
        # so we use hi-level locking
        with task_start_lock:
            cmd = self.app.executor + [self.app.binpath]
            self.proc = subprocess.Popen(cmd, env=self.env, cwd=self.cwd, 
                                     stdout=self.stdout_file, stderr=self.stderr_file)
        
        salsa.core.logger.log(LoggerSeverity.LS_INFO, 'Started task %s' % str(self))
        
        self.proc_retcode = self.proc.wait()
        
        salsa.core.logger.log(LoggerSeverity.LS_INFO, 'Finished task %s, return code: %d' % (self.taskid, self.proc_retcode))
    
    def try_finish(self):
        self.proc.send_signal(signal.SIGINT)
            
        finished = self.proc_finish_event.wait(TASK_INTR_WAIT)
    
        if not finished:
            salsa.core.logger.log(LoggerSeverity.LS_INFO, 'Timed out waiting for task %s finish, retry with SIGKILL' % self.taskid)
            self.proc.kill()
    
    def finish(self):
        # Task waits TASK_INTR_WAIT seconds for subprocess to finish after SIGINT
        # otherwise retries with SIGKILL
        if self.proc:
            if trace_exec_flow: do_trace('Task %s is finishing' % self.taskid)
            
            self.finish_flag = True
            self.keep_alive = False
            
            try:
                self.try_finish()
            except OSError:
                pass
            
        
class Service(Task):
    def __init__(self, app, port):
        Task.__init__(self, app)
        
        self.port = port
        
        self.set_keep_alive()
    
    def svc_info(self):
        return {'name' : self.app.name, 'port': self.port, 'host': 'localhost'}

class TaskManager:
    def __init__(self):
        self.apps = []
        
        self.tasks = []
        self.task_lock = threading.RLock()
        
        self.disp_queue = Queue()
        
        self.out_dir =  salsa.core.abs_data_path('out/', cleanup = True)
        
        
        self.make_task_env()
        self.config()
    
    def config(self):
        self.max_exec_pool = int(salsa.core.config.get_global_var('max_exec_pool'))
        
        apps_node = salsa.core.config.get_section('applications')
        apps_cfg = salsa.core.config.process_list(apps_node, 'application', 
                                                   ['name', 'binpath'], 
                                                   {'type': 'app', 'exec_name': ''}, 
                                                   need_el_ref=True)
        
        for app_cfg in apps_cfg:
            try:
                self.create_app(**app_cfg)
            except TypeError as te:
                app_name = app_cfg.get('name', '<unknown>')
                
                raise TaskMgrException('Error in <applications> section in application declaration %s' % app_name)
        
    def make_task_env(self):
        def append_path(env, k, x): 
            env[k] = os.pathsep.join(env.get(k, '').split(os.pathsep) + [x])
            
        env = os.environ.copy() 
        env['SALSA_CORE_PORT'] = str(salsa.core.core_port)
        env['SALSA_HOST'] = salsa.core.core_host
        
        append_path(env, 'PATH', salsa.core.abs_bin_path('bin/'))
        append_path(env, 'PYTHONPATH', salsa.core.abs_bin_path('lib/'))
        
        self.default_task_env = env
    
    def create_app(self, name, binpath, type, exec_name, _el):
        params = salsa.core.config.process_parameters(_el)
        
        app = Application(name)
        app.set_binpath(binpath)
        app.set_params(params)
        app.set_exec(exec_name)        
        
        if type == 'service':
            self.config_service(app, _el)
        
        self.apps.append(app)
        salsa.core.logger.log(LoggerSeverity.LS_DEBUG, 'Registered application %s @%s' % (app.name, app.pretty_binpath))
    
    def config_service(self, app, _el):
        svc_node = salsa.core.config.try_find(_el, 'service')
        port = int(svc_node.attrib['port'])
        
        self.create_service(app, port)
    
    def create_service(self, app, port):
        # Check if any service with same name and port exists 
        if any((svc.app.name == app.name or svc.port == port) 
               for svc in self.tasks 
               if isinstance(svc, Service)):
            raise TaskMgrException('Service named %s duplicated in config or port is already in use')
        
        # Create service and run
        
        svc = Service(app, port)
        self.tasks.append(svc)
        
        salsa.core.logger.log(LoggerSeverity.LS_DEBUG, 'Registered service %s @%s port: %d' % 
                              (app.name, app.pretty_binpath, svc.port))
        
        self.init_task(svc, {})
    
    def start(self):
        for svc in self.tasks:
            svc.start(self.finish_task_handler)
    
    def finish_task_handler(self, task):
        with self.task_lock:
            self.tasks.remove(task)
            
            if not self.disp_queue.empty() and self.is_slots_available():
                (reply, task) = self.disp_queue.get()
                
                task.start(self.finish_task_handler)
                
                reply.wakeup()
    
    def find_task(self, taskid):
        for task in self.tasks:
            if task.taskid == taskid:
                return task
        
        raise KeyError(taskid)
    
    def stop(self):
        # self.tasks is mutable because of finish_task_handler
        # so, make a copy
        tasklist = self.tasks[:]
        
        for task in tasklist:
            task.finish()
    
    def is_slots_available(self):
        slots_busy = len(filter(lambda task: not isinstance(task, Service), self.tasks))
        
        return slots_busy < self.max_exec_pool
    
    def attach_tasks(self, parentid, tasks):
        with self.task_lock:
            attachable_tasks = []
            
            # Not all tasks are attached because they may be already finished
            for task in self.tasks:
                if task.taskid in tasks:
                    attachable_tasks.append(task)
            
            ignored_taskids = list(set(tasks) - set(task.taskid for task in attachable_tasks))
            salsa.core.logger.log(LoggerSeverity.LS_DEBUG, 'Task %s not attaches to tasks %s' % (parentid, ','.join(ignored_taskids)))            
            
            reply = IODeferredReply({}, TaskWaiterEvent(len(attachable_tasks)))
            
            for task in attachable_tasks:                
                task.attach(reply)
            
            salsa.core.logger.log(LoggerSeverity.LS_DEBUG, 'Task %s attached to %s' % 
                                    (parentid, ', '.join(task.taskid for task in attachable_tasks)))    
            
            return reply
    
    def init_task(self, task, args):
        task.set_env(self.default_task_env.copy())
        task.set_dirs(salsa.core.abs_data_path(), self.out_dir)
        task.set_args(args)
    
    def create_task(self, name, args = {}, parentid = ''):
        salsa.core.logger.log(LoggerSeverity.LS_DEBUG, 'Trying to create task %s parent: %s' % (name, parentid))
               
        for app in self.apps:
            if app.name == name:
                break
        else:
            raise RPCError(RPCError.BAD_REQUEST, 'Application %s not found' % name)
        
        task = Task(app)
        self.init_task(task, args)
        
        with self.task_lock:
                self.tasks.append(task)
        
        if self.is_slots_available():            
            task.start(self.finish_task_handler)
        
            return task.taskid
        else:
            salsa.core.logger.log(LoggerSeverity.LS_INFO, 'Task %s start is delayed' % (task.taskid))
            
            reply = IODeferredReply(task.taskid, threading.Event())
            
            self.disp_queue.put((reply, task))
            
            return reply
    
    def get_task_info(self, taskid):
        task = self.find_task(taskid)
        info = {'args': task.args, 'type': 'task', 
                'name': task.app.name, 
                'params': task.app.params,
                'pid': task.get_pid()}
        
        if not task:
            raise RPCError(RPCError.BAD_REQUEST, 'Task %s not found' % taskid)
        
        if isinstance(task, Service):
            info['type'] = 'service'
            
            info.update(task.svc_info())
            
        return info
    
    def list_services(self):
        with self.task_lock:
            return IOFeederReply([svc.svc_info() 
                                  for svc in self.tasks
                                  if isinstance(svc, Service)])
    
    def list_tasks(self):
        with self.task_lock:
            return IOFeederReply([task.taskid
                                  for task in self.tasks])
